from torch import nn
import torch
class TensorReduction(nn.Module):
    def __init__(self, dim1, dim2):
        super(TensorReduction, self).__init__()
        # 定义一个可训练的权重参数，维度为(dim2, dim1, dim1)
        self.weight = nn.Parameter(torch.rand(dim2, dim1, dim1))

    def forward(self, X):
        # 初始化一个全零张量，大小为(X.shape[0], self.weight.shape[0])
        Y = torch.zeros(X.shape[0], self.weight.shape[0])
        for k in range(self.weight.shape[0]):
            # 计算temp = X @ weight[k] @ X^T
            temp = torch.matmul(X, self.weight[k]) @ X.T
            # 取temp的对角线元素，存入Y[:, k]
            Y[:, k] = temp.diagonal()
        return Y

# 创建一个TensorReduction层，dim1=10, dim2=5
layer = TensorReduction(10, 5)
# 创建一个大小为(2, 10)的张量X
X = torch.rand(2, 10)
# 对layer(X)进行前向传播，返回一个大小为(2, 5)的张量
print(layer(X).shape)